package org.xyl.engine;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;

public class SimpleRAG {
    public static void main(String[] args) {
        // 1. 创建嵌入模型
        EmbeddingModel embeddingModel = new dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel();

        // 2. 创建内存向量存储
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        // 3. 添加文档
        String document = "Java是一种面向对象的编程语言，由Sun Microsystems开发。";
        TextSegment segment = TextSegment.from(document);
        Embedding embedding = embeddingModel.embed(document).content();
        embeddingStore.add(embedding, segment);

        // 4. 简单查询示例
        String query = "Java有什么特点？";
        Embedding queryEmbedding = embeddingModel.embed(query).content();
        
        // 使用嵌入存储直接查询相关内容
        List<?> results = embeddingStore.findRelevant(queryEmbedding, 1);
        
        // 打印结果
        System.out.println("查询：" + query);
        System.out.println("找到的相关结果数：" + results.size());
        System.out.println("示例RAG流程完成");
    }
}